# -*- coding: utf-8 -*-

__author__="Boris Avdeev"
__date__ ="$Mar 8, 2011 10:09:48 AM$"

from TCdata import TCdata
import sqlite3 as sql
import numpy as np
import pymc as pm

class UTHAgeModel(object):
    """Single grain and sample age model (identicsl to the generic AgeModel at the moment)"""
    def __init__(self, uth, sample_age, RelErr=None, Age=None):
        """If Age or RelErr are None, they are modelled. Alternatively, constants or extermal models can be passed."""
#        tries = 1000
        max_age=4500.
        name=uth.name
        #eU = uth.eU
        #Ft = uth.Ft
        #He = uth.He
        ages = uth.farley_ages
#        while tries > 0:
#            try:
        # PRIORS
        if RelErr is None:
            if sample_age: self.RelErr = pm.Uniform("AHeRelErr_%s"%name, 0.05, 0.5, value=0.1)
            else: self.RelErr = 0.15  # Do not model error for single grain ages: no such information! 15% is the Caucasus average.
        #elif isinstance(RelErr,float):
        #    self.RelErr =
        else: self.RelErr = RelErr
        #self.eU =    pm.Uniform("eU_%s"    % name, 0., 10., value = eU, plot=False)

        if Age is None:
            if sample_age: # assume common age for all grains
                init_age = np.mean(ages)
            else: # each grain has individual age
                init_age = ages
            self.Age = pm.Uniform("AHeAge_%s"  % name, 0., max_age, value = init_age)
        else: self.Age = Age

        #TODO: DATA
        self.ObsAge = pm.Normal("ObsAHeAge_%s" % name,  mu = self.Age, tau = 1./(self.RelErr*self.Age)**2, value = ages, observed=True)
        #self.ObseU = pm.Normal("ObseU_%s" % name,  mu = self.P, tau = 1./self.P1s**2, value = P, observed=True)
        #self.ExpHe = pm.Lambda("ExpNs_%s"%name,
        #    lambda t=self.Age, a=A, p=self.P, zeta=self.Zeta, g=uth.g : uth_exp_ns(t,a,p,zeta,g),plot=False)
        #self.ObsHe = pm.Poisson("Ns_%s"%name, mu = self.ExpNs, value = Ns, observed=True)

        # SIMULATION
        #self.SimNs = pm.Deterministic(eval=pm.rpoisson,name="SimNs_%s"%name,doc='', parents={'mu':self.ExpNs},plot=False)
        self.SimAge = pm.Lambda("SimAHeAge_%s"%name,lambda mu=self.Age,re=self.RelErr: pm.rnormal(mu,1./(re*mu)**2,size=len(ages)),plot=False)
 #               break
 #           except pm.ZeroProbability:
 #               tries-=1
 #       if tries==0: raise pm.ZeroProbability

    def get_list(self):
        return self.__dict__.values()




class UTHdata(TCdata):
    def __init__(self, sample):
        TCdata.__init__(self)
        conn = sql.connect('/raid/work1/Caucasus/db/data.sqlite')
        query = "SELECT sample,age,U,Th,Sm,He,mass,Ft,r,l FROM ahe_sg WHERE sample=?"
        table = conn.cursor().execute(query, (sample,)).fetchall()
        conn.close()
        if len(table)==0: raise KeyError("No UTH data found for sample %s"%sample)
        table = np.array(table)
        self.farley_ages =  table[:,1].astype(float)#[row[1] for row in table]
        self.quick_ages = self.farley_ages  # a plug
        self.method = 'UTH'
        self.name = sample
        self.n_cnt = len(self.farley_ages)

        #-------INITIATE MODELS-------------------------------------------------
        try: self.sa_model = UTHAgeModel(self, True)
        except pm.ZeroProbability: print "Cannot initiate SA model"
        try: self.sga_model = UTHAgeModel(self, False)
        except pm.ZeroProbability: print "Cannot initiate SGA model"
